# build_check.py
# usage: (第一遍SQL生成后）根据tmp_question2sql.json已经归集的问题大类分组，生成人工友好的待标记清单
# python ./src/tools/build_check.py --questions "./data/questions.round02.jsonl" --results "./data/samples.round02.jsonl" --q2s "./tmp_question2sql.json" --check "./tmp_checklist.txt"

import argparse
import json
from collections import defaultdict
from typing import Dict, List, Any

# 假设 clean_question 定义在 clean_question.py 中
from clean_question import clean_question


def load_jsonl(file_path: str) -> Dict[int, Any]:
    """
    加载 JSONL 文件，返回以 id 为键的字典。
    """
    data = {}
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            item = json.loads(line)
            data[item['id']] = item
    return data


def load_q2s(file_path: str) -> Dict[str, str]:
    """
    加载 q2s 文件，返回一个字典：short_question -> table
    注意：一个 short_question 只能属于一个 table（按首次出现为准）
    """
    q2s_map = {}
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            item = json.loads(line)
            tables = item.get('table', [])
            questions = item.get('questions', [])
            for q in questions:
                if q not in q2s_map:
                    q2s_map[q] = tables[0] if tables else "unknown_table"
    return q2s_map


def main():
    parser = argparse.ArgumentParser(description="Build check file by mapping questions to tables via short_question.")
    parser.add_argument('--questions', required=True, help='Path to questions JSONL file')
    parser.add_argument('--results', required=True, help='Path to results JSONL file')
    parser.add_argument('--q2s', required=True, help='Path to q2s mapping JSONL file')
    parser.add_argument('--check', required=True, help='Path to output check file')

    args = parser.parse_args()

    # 1. 加载数据
    print("Loading questions...")
    questions = load_jsonl(args.questions)  # id -> {id, question, level}

    print("Loading results...")
    results = load_jsonl(args.results)     # id -> {id, sql}

    print("Loading q2s mapping...")
    q2s_map = load_q2s(args.q2s)          # short_question -> table

    # 2. 构建输出结构：table -> short_question -> list of entries
    output_data = defaultdict(lambda: defaultdict(list))

    # 3. 遍历 results 文件中所有记录
    for result_id, result_item in results.items():
        if result_id not in questions:
            print(f"Warning: result id={result_id} not found in questions, skipped.")
            continue

        question_item = questions[result_id]
        raw_question = question_item['question']
        level = question_item['level']

        # 调用 clean_question 获取标准化短问题
        try:
            short_q = clean_question(raw_question)
        except Exception as e:
            print(f"Error cleaning question id={result_id}: {e}")
            continue

        # 查找对应的表
        if short_q not in q2s_map:
            print(f"Warning: short_question '{short_q}' not found in q2s mapping, skipped.")
            continue

        table = q2s_map[short_q]

        # 构造输出条目
        entry = {
            "id": result_id,
            "level": level,
            "flag": 0,
            "question": raw_question,
            "sql": result_item["sql"]
        }

        output_data[table][short_q].append(entry)

    # 4. 按格式写入 check 文件
    with open(args.check, 'w', encoding='utf-8') as f:
        for table in sorted(output_data.keys()):
            f.write(f"{table}\n")
            short_questions = output_data[table]
            for short_q in sorted(short_questions.keys()):
                f.write(f"    {short_q}\n")
                for entry in short_questions[short_q]:
                    # 使用 json.dumps 美化输出，但保持单行
                    f.write(f"        {json.dumps(entry, ensure_ascii=False)}\n")

    print(f"Check file generated at: {args.check}")


if __name__ == "__main__":
    main()